import os
import numpy as np

import lpips
from tqdm import tqdm
from torch_fidelity import calculate_metrics


##############################
## 计算FID值（自然真实程度）######
##############################
def cal_fid(adv_path, ori_path):
    # Each input can be either a string (path to images, registered input), or a Dataset instance
    metrics_dict = calculate_metrics(adv_path, ori_path, cuda=True, isc=False, fid=True, kid=False, verbose=False)
    fid = metrics_dict['frechet_inception_distance']
    score_fid = np.sqrt(1 - np.min([fid, 200]) / 200)
    print('score_fid:', score_fid)
    return score_fid

##############################
## 计算LPIPS值(和原图的感知距离)##
##############################
def cal_lpips(adv_path, ori_path):
    lpips_out_dir = './lpips_out.txt'  # 输出文件
    # 初始化模型
    loss_fn = lpips.LPIPS(net='vgg').cuda()
    # 待比较文件序列
    f = open(lpips_out_dir,'w')
    files_list = os.listdir(adv_path)
    files_list.sort(key=lambda x:int(x.split('.')[0]))  # 排序文件名
    lpips_dis = []
    # 开始比较
    for file in tqdm(files_list):
        
        if(os.path.exists(os.path.join(ori_path,file))):
            # Load images
            img0 = lpips.im2tensor(lpips.load_image(os.path.join(adv_path,file))).cuda() # RGB image from [-1,1]
            img1 = lpips.im2tensor(lpips.load_image(os.path.join(ori_path,file))).cuda()

            # Compute distance
            dist01 = loss_fn.forward(img0,img1)
            # print('%s: %.3f'%(file,dist01))
            f.writelines('%s: %.6f\n'%(file,dist01))
            lpips_dis.append(dist01.cpu().detach().numpy().max())
    f.close()
    score_lpips = np.mean(lpips_dis)
    score_lpips = np.sqrt( 1 - 2 * (np.min([np.max([0.2, score_lpips]), 0.7]) - 0.2) )
    print('score_lpips: ', score_lpips)
    return score_lpips
